package com.k8.common.executor.threadPool.extendImpl;

import com.k8.common.util.DoublyLinkedList;
import org.jetbrains.annotations.NotNull;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 阻塞队列同lock中aqs的s相同，都是通过lockSupport的park，unpark进行阻塞和唤醒
 * <p>
 * 线程池getTask时通过take获取，会阻塞，在SynchronousQueue中会unpark住，
 * 而提交任务，offer如果设置时间，则会park该时间，不设置则直接park(0)立即返回
 * 所以有一点概率在增加线程时将任务提交给了恰好执行完毕的线程导致不会新建线程，而如果此时线程都忙则super.execute中
 * offer不会被空闲线程的take获取，则会创建新线程，所以在提交任务时一定要确保等待队列里有任务，如此才可以确保线程都忙，此外
 * 为了防止提交的任务中
 * <p>
 * 支持检测超时任务，当有线程超时时可以动态增加线程
 * <p>
 * 给父类一个SynchronousQueue，以实现零缓存的阻塞队列
 * 自身还需要维护一个阻塞队列，一般是无界队列等有一定缓存能力的
 * <p>
 * ****主要通过两个守护线程来维护任务从自身任务队列到父类的任务队列的transfer，直接在阻塞队列中转移可以使得线程的增加由派生类控制
 * 即负责超时增加线程的守护线程来进行增加线程
 */
public class SmartDynamicThreadPoolImpl extends ThreadPoolExecutor {
    private final BlockingQueue<Runnable> workQueue;
    private final BlockingQueue<Runnable> superWorkQueue;
    private final DoublyLinkedList<TaskContext> submittedTasks = new DoublyLinkedList<>();
    private Semaphore timeOutSemaphore = new Semaphore(0);
    private Semaphore shutDownSemaphore = new Semaphore(0);
    private final Lock preSubmittedTasksLock = new ReentrantLock();
    private final Lock lastSubmittedTasksLock = new ReentrantLock();
    //毫秒
    private final long timeOut;

    private volatile Thread transferThread;
    private volatile Thread outTimeAddThread;
    private final Lock mainLock = new ReentrantLock();
    private volatile AtomicBoolean isAvailable = new AtomicBoolean(true);//可用的
    private volatile boolean shutDown = false;
    private volatile boolean shutDownNow = false;
    private volatile AtomicInteger taskCount = new AtomicInteger(0);
    //用于保证守护线程启动成功后其他的投递才生效
    private volatile boolean startSuccess0 = false;
    private volatile boolean startSuccess1 = false;

    private final RejectedExecutionHandler rejectHandler;

    public SmartDynamicThreadPoolImpl(long timeOut, int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize + 2, maximumPoolSize + 2, keepAliveTime, unit, new SynchronousQueue<>(), new RejectHandler(handler));
        if (timeOut <= 0) throw new IllegalArgumentException("TimeOut can not less than 0");
        this.superWorkQueue = super.getQueue();
        this.workQueue = workQueue;
        this.timeOut = timeOut;
        this.rejectHandler = handler;
        startWatchDog();
    }

    public SmartDynamicThreadPoolImpl(long timeOut, int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull ThreadFactory threadFactory, @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize + 2, maximumPoolSize + 2, keepAliveTime, unit, new SynchronousQueue<>(), threadFactory, new RejectHandler(handler));
        if (timeOut <= 0) throw new IllegalArgumentException("TimeOut can not less than 0");
        this.superWorkQueue = super.getQueue();
        this.workQueue = workQueue;
        this.timeOut = timeOut;
        this.rejectHandler = handler;
        startWatchDog();
    }

    public SmartDynamicThreadPoolImpl(long timeOut, int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue) {
        super(corePoolSize + 2, maximumPoolSize + 2, keepAliveTime, unit, new SynchronousQueue<>(), new RejectHandler(new AbortPolicy()));
        if (timeOut <= 0) throw new IllegalArgumentException("TimeOut can not less than 0");
        this.superWorkQueue = super.getQueue();
        this.workQueue = workQueue;
        this.timeOut = timeOut;
        this.rejectHandler = ((RejectHandler) super.getRejectedExecutionHandler()).getRejectedExecutionHandler();
        startWatchDog();
    }

    public SmartDynamicThreadPoolImpl(long timeOut, int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull ThreadFactory threadFactory) {
        super(corePoolSize + 2, maximumPoolSize + 2, keepAliveTime, unit, new SynchronousQueue<>(), threadFactory, new RejectHandler(new AbortPolicy()));
        if (timeOut <= 0) throw new IllegalArgumentException("TimeOut can not less than 0");
        this.superWorkQueue = super.getQueue();
        this.workQueue = workQueue;
        this.timeOut = timeOut;
        this.rejectHandler = ((RejectHandler) super.getRejectedExecutionHandler()).getRejectedExecutionHandler();
        startWatchDog();
    }


    public static class TaskContext {
        private final long startTime;
        public static final ThreadLocal<TaskContext> taskContextLocal = new ThreadLocal<>();
        private volatile boolean end;
        //记录该任务上一次被统计过的超时了多长时间，下一次再统计的时候要去除这一部分
        private volatile long outTime;
        private final Thread thread;

        public TaskContext(long startTime, Thread thread) {
            this.startTime = startTime;
            this.thread = thread;
        }

        public boolean isEnd() {
            return end;
        }

        public boolean isCanRun() {
            return !this.thread.isInterrupted() && !this.thread.isAlive();
        }

        public void setEnd(boolean end) {
            this.end = end;
        }
    }

    @Override
    public BlockingQueue<Runnable> getQueue() {
        return this.workQueue;
    }

    @Override
    public RejectedExecutionHandler getRejectedExecutionHandler() {
        return this.rejectHandler;
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        if (!(r instanceof AddRunnable)) {
            TaskContext taskContext = null;
            try {
                lastSubmittedTasksLock.lock();
                taskContext = new TaskContext(System.currentTimeMillis(), Thread.currentThread());
                //添加进计时队列
                submittedTasks.addLast(taskContext);
            } finally {
                lastSubmittedTasksLock.unlock();
                TaskContext.taskContextLocal.set(taskContext);
                taskCount.incrementAndGet();
            }
        }
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        if (!(r instanceof AddRunnable)) {
            TaskContext taskContext = TaskContext.taskContextLocal.get();
            taskContext.setEnd(true);
            TaskContext.taskContextLocal.remove();
            boolean lockSuccess = false;
            try {
                lockSuccess = preSubmittedTasksLock.tryLock();
                if (lockSuccess) {
                    int size = submittedTasks.size();
                    if (size > 1) {
                        DoublyLinkedList.Node<TaskContext> node = submittedTasks.getNodeAt(0);
                        for (int i = 1; i < size; ++i) {
                            TaskContext taskC = node.getValue();
                            if (taskC.isEnd() || !taskC.isCanRun()) {
                                submittedTasks.remove(node);
                            }
                            node = node.next;
                        }
                    }
                }
            } finally {
                if (lockSuccess) {
                    preSubmittedTasksLock.unlock();
                }
                taskCount.decrementAndGet();
            }
        }
    }


    @Override
    public void execute(@NotNull Runnable command) {
        if (!isAvailable.get()) {
            rejectHandler.rejectedExecution(command, this);
            return;
        }
        while (!startSuccess0 || !startSuccess1) {
        }
        if (!coreExecute(command)) {
            if (!this.workQueue.offer(command)) {
                rejectHandler.rejectedExecution(command, this);
            }
        }
    }

    private boolean coreExecute(Runnable command) {
        if (super.getPoolSize() < super.getCorePoolSize()) {
            try {
                mainLock.lock();
                if (super.getPoolSize() < super.getCorePoolSize()) {
                    super.execute(command);
                    return true;
                }
                return false;
            } finally {
                mainLock.unlock();
            }
        }
        return false;
    }

    /**
     * 先设置线程池不可用标志，拒绝接收新任务
     * 两个守护线程也会因标志位的变化而中断
     * 清空已提交任务的上下文
     * 再执行父类的shutdownNow，线程不再getTask消耗次级workQueue中的任务，且修改标志位且中断所有线程
     * 最后执行父类的shutDown
     */
    @Override
    public List<Runnable> shutdownNow() {
        if (isAvailable.get()) {
            if (!isAvailable.compareAndSet(true, false)) Collections.emptyList();
        } else {
            return Collections.emptyList();
        }
        shutDownNow = true;
        try {
            shutDownSemaphore.acquire(2);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        List<Runnable> runnables = super.shutdownNow();
        List<Runnable> list = this.workQueue.stream().toList();
        if (runnables.size() > 0) {
            list.addAll(runnables);
        }
        return list;
    }

    /**
     * 先设置线程池不可用标志，拒绝接收新任务，transferThread并不会立即中断，而是处理完剩余所有剩余任务通过现有线程
     * shutDown不会将标志位修改为stop，故此任务仍然需要执行，所以此时需执行 super.shutdown(); 修改标志位且中断所有线程
     * 最后执行父类的shutDown
     */
    @Override
    public void shutdown() {
        if (isAvailable.get()) {
            if (!isAvailable.compareAndSet(true, false)) return;
        } else {
            return;
        }
        shutDown = true;
        try {
            shutDownSemaphore.acquire(2);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        super.shutdown();
    }

    public boolean isAvailable() {
        return this.isAvailable.get();
    }

    /**
     * submit最终还是调用execute，并不需要重写
     * */
    /*@NotNull
    @Override
    public <T> Future<T> submit(@NotNull Runnable task, T result) {
        if (task == null) throw new NullPointerException();
        RunnableFuture<T> ftask = newTaskFor(task, result);
        execute(ftask);
        return ftask;
    }*/

    /**
     * invoke风格的提交不用进行重写，走码发现最终还是会调用到execute，使用本逻辑
     */
    /*@Override
    public <T> T invokeAny(@NotNull Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
        return super.invokeAny(tasks);
    }
    @NotNull
    @Override
    public <T> List<Future<T>> invokeAll(@NotNull Collection<? extends Callable<T>> tasks, long timeout, @NotNull TimeUnit unit) throws InterruptedException {
        return super.invokeAll(tasks, timeout, unit);
    }*/
    @Override
    public long getTaskCount() {
        return this.taskCount.get();
    }

    private void startWatchDog() {
        AddRunnable transfer = () -> {
            startSuccess0 = true;
            transferThread = Thread.currentThread();
            try {
                Runnable task = null;
                while (isAvailable.get() || (shutDown && !workQueue.isEmpty())) {
                    task = workQueue.take();
                    if (task == null) continue;
                    boolean offer = superWorkQueue.offer(task, timeOut, TimeUnit.MILLISECONDS);
                    while (!offer && (isAvailable.get() || shutDown)) {
                        //超时了仍然没有线程来处理该任务
                        //提交给另外一个专门处理超时的线程
                        //此时若已经开始停止了，就不再申请扩容了
                        if (isAvailable.get()) {
                            timeOutSemaphore.release();
                        }
                        offer = superWorkQueue.offer(task, timeOut, TimeUnit.MILLISECONDS);
                    }
                    if (shutDownNow && !offer) {
                        //如果在task取出之后，workQueue又添加进来一个，此时达到队列最大值，然后shutDownNow，那么就会offer失败
                        if (!workQueue.offer(task)) {
                            rejectHandler.rejectedExecution(task, this);
                        }
                    }
                }
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            shutDownSemaphore.release();
            transferThread.interrupt();
        };
        AddRunnable timeOutResolve = () -> {
            startSuccess1 = true;
            outTimeAddThread = Thread.currentThread();
            int count = 0;
            while (isAvailable.get()) {
                try {
                    boolean acquire = timeOutSemaphore.tryAcquire(timeOut, TimeUnit.MILLISECONDS);
                    int maxAddThreads = super.getMaximumPoolSize() - super.getPoolSize();
                    if (acquire) {
                        //如果出现这种情况，则说明当前线程池所有线程全部超时，直接扩容1.5
                        int addCount = Math.min(Math.max((int) (super.getPoolSize() * 0.5), 1), maxAddThreads);
                        addThread(addCount);
                    } else {
                        //自检超时任务
                        preSubmittedTasksLock.lock();
                        int size = submittedTasks.size();
                        if (size > 1) {
                            try {
                                DoublyLinkedList.Node<TaskContext> node = submittedTasks.getNodeAt(0);
                                long allOutTime = 0;
                                for (int i = 1; i < size; ++i) {
                                    TaskContext taskC = node.getValue();
                                    if (taskC.isEnd() || !taskC.isCanRun()) {
                                        submittedTasks.remove(node);
                                    } else {
                                        long runTime = System.currentTimeMillis() - taskC.startTime;
                                        if (runTime > timeOut) {
                                            long lastOutTime = taskC.outTime;
                                            //为了避免某一个线程已经运行了很长时间，所以在计算它的超时权值时，最大值为timeOut
                                            long addOutTime = Math.min(runTime - timeOut - lastOutTime, timeOut);
                                            taskC.outTime = runTime - timeOut;
                                            allOutTime += addOutTime;
                                        }
                                    }
                                    node = node.next;
                                }
                                //超时总时长共为多少个timeOut，如果此时未超时线程个数大于工作队列中的线程数，那么未来是有可能成功的
                                //设置一个合理的计算公式，根据noOutTimeTasks
                                int timeOutCount = (int) (allOutTime / timeOut);
                                int candidateTasks = workQueue.size() + 1;
                                int addCount = Math.min(timeOutCount, candidateTasks);
                                addCount = Math.min(addCount, maxAddThreads);
                                addThread(addCount);
                            } finally {
                                preSubmittedTasksLock.unlock();
                            }
                        }
                    }
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
            shutDownSemaphore.release();
            outTimeAddThread.interrupt();
        };

        super.execute(transfer);
        super.execute(timeOutResolve);
    }

    private void addThread(int count) {
        int poolSize = super.getPoolSize();
        //理论上应该还要+1，但是后续区分普通任务和这个被弹出队列的任务不太方便，即不考虑，而一个任务的影响也不大
        int workSize = workQueue.size();

        for (int i = 0; i < count; ++i) {
            if (workSize == 0 || !isAvailable.get()) return;
            super.execute(new AddRunnable() {
                @Override
                public void run() {
                }
            });
            for (; ; ) {
                if (!isAvailable.get()) break;
                if (poolSize == super.getPoolSize()) {
                    //1.被super突然空闲的线程给take走，有空闲线程了，则等待工作队列的任务数减1
                    //2.有线程被非shutDown中断，但是提交成功，等待队列中任务减少
                    //空任务执行很快
                    while (workSize == workQueue.size() && isAvailable.get()) {
                        Thread.onSpinWait();
                    }
                    workSize = workQueue.size();
                } else if (poolSize < super.getPoolSize()) {
                    //添加线程成功
                    poolSize = super.getPoolSize();
                    break;
                } else {
                    //反而线程数减小，则说明有线程shutDown被中断，则添加需要增加的线程数
                    count += Math.max(0, poolSize - super.getPoolSize());
                    poolSize = super.getPoolSize();
                }
                if (workSize == 0) return;
                //还是有任务在工作队列,再次尝试增加线程
                super.execute(new AddRunnable() {
                    @Override
                    public void run() {
                    }
                });
            }
        }
    }

    private interface AddRunnable extends Runnable {

    }

    private static class RejectHandler implements RejectedExecutionHandler {
        private final RejectedExecutionHandler rejectedExecutionHandler;

        public RejectedExecutionHandler getRejectedExecutionHandler() {
            return rejectedExecutionHandler;
        }

        public RejectHandler(RejectedExecutionHandler rejectedExecutionHandler) {
            this.rejectedExecutionHandler = rejectedExecutionHandler;
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            this.rejectedExecutionHandler.rejectedExecution(r, executor);
        }
    }
}
